import cv2#调用cv2库
import os#调用os库

# 加载级联分类器（确保文件路径正确）
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

# 打开摄像头
cap = cv2.VideoCapture(0)

# 数据集路径
dataset_path = 'dataset'

# 创建数据集文件夹（如果不存在）
if not os.path.exists(dataset_path):
    os.makedirs(dataset_path)

# 输入用户 ID
user_id = input("请输入用户 ID（用于创建文件夹）：")

# 用户文件夹路径
user_folder = os.path.join(dataset_path, user_id)
if not os.path.exists(user_folder):
    os.makedirs(user_folder)

# 照片计数器
photo_count = 0

while True:
    # 读取一帧图像
    ret, frame = cap.read()
    if not ret:
        break

    # 转换为灰度图（提高检测效率）
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 人脸检测（调整参数优化效果）
    faces = face_cascade.detectMultiScale(
        gray,
        scaleFactor=1.1,
        minNeighbors=5,
        minSize=(30, 30)
    )

    # 在图像上绘制矩形框
    for (x, y, w, h) in faces:
        cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)

    # 显示结果
    cv2.imshow('Face Detection', frame)

    # 按下 's' 键拍照
    key = cv2.waitKey(1)
    if key == ord('s'):
        if len(faces) > 0:  # 检测到人脸时拍照
            photo_count += 1
            photo_path = os.path.join(user_folder, f"{user_id}_{photo_count}.jpg")
            cv2.imwrite(photo_path, frame)
            print(f"拍照成功：{photo_path}")
        else:
            print("未检测到人脸，无法拍照")

    # 按 'q' 退出循环
    if key == ord('q'):
        break

# 释放资源
cap.release()
cv2.destroyAllWindows()